// Copyright (C) 2018-2023 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#include "single_layer_tests/convolution.hpp"

#include <vector>

#include "../skip_tests_check.hpp"
#include "common_test_utils/test_constants.hpp"

using namespace LayerTestsDefinitions;

namespace {

class gna_convolution_layerTest : public ConvolutionLayerTest {
protected:
    void Run() override {
        ConvolutionLayerTest::Run();
    }

    void SetUp() override {
        ConvolutionLayerTest::SetUp();
    }
};

TEST_P(gna_convolution_layerTest, CompareWithRefs) {
    Run();
}

const std::vector<InferenceEngine::Precision> netPrecisions = {InferenceEngine::Precision::FP32,
                                                               InferenceEngine::Precision::FP16};

/* ============= 2D Convolution ============= */
const std::vector<std::vector<size_t>> kernelsH1 = {{1, 3}, {1, 5}};
const std::vector<std::vector<size_t>> stridesH1 = {{1, 1}, {1, 3}};
const std::vector<std::vector<size_t>> stridesH1Big = {{1, 9}, {1, 16}};
const std::vector<std::vector<ptrdiff_t>> padBeginsH1 = {{1, 0}, {1, 3}};
const std::vector<std::vector<ptrdiff_t>> padEndsH1 = {{1, 0}, {1, 3}};
const std::vector<std::vector<size_t>> dilationsH1 = {{1, 1}, {1, 3}};
const std::vector<std::vector<size_t>> inputShapesH1 = {{1, 1, 1, 32}, {1, 32, 1, 160}, {1, 8, 1, 64}};
const std::vector<std::vector<size_t>> kernelsW1 = {{3, 1}, {5, 1}};
const std::vector<std::vector<size_t>> stridesW1 = {{1, 1}, {3, 1}};
const std::vector<std::vector<ptrdiff_t>> padBeginsW1 = {{0, 1}, {3, 1}};
const std::vector<std::vector<ptrdiff_t>> padEndsW1 = {{0, 1}, {3, 1}};
const std::vector<std::vector<size_t>> dilationsW1 = {{1, 1}, {3, 1}};
const std::vector<std::vector<size_t>> inputShapesW1 = {{1, 1, 32, 1}, {1, 32, 160, 1}, {1, 8, 64, 1}};
const std::vector<size_t> numOutCannels = {4, 8, 12};

const std::vector<std::vector<size_t>> kernels2D = {
    {5, 1},
    {4, 1},
    {1, 3},
    {1, 2},
    {2, 2},
    {7, 1},
    {3, 3},
};

const std::vector<std::vector<size_t>> kernels2D_big = {
    {7, 2},
    {2, 7},
    {3, 7},
    {6, 6},
    {7, 7},
};

const std::vector<std::vector<size_t>> kernels2D_3x3 = {
    {3, 3},
};
const std::vector<std::vector<size_t>> kernels2D_5x6 = {
    {5, 6},
};

const std::vector<std::vector<size_t>> strides2D = {
    {1, 1},
};
const std::vector<std::vector<ptrdiff_t>> padBegins2D = {
    {0, 0},
};
const std::vector<std::vector<ptrdiff_t>> padEnds2D = {
    {0, 0},
};
const std::vector<std::vector<size_t>> dilations2D = {
    {1, 1},
};
const std::vector<size_t> numOutCannels2D = {8, 16, 32};

const std::vector<size_t> num_out_channels_for_mapped_2d = {4, 8, 12};

const std::vector<size_t> input2DNCHW = {1, 8, 20, 16};

const std::vector<size_t> input2DNCHW_3x3 = {1, 16, 3, 3};
const std::vector<size_t> input2DNCHW_5x6 = {1, 16, 5, 6};

const std::vector<std::vector<size_t>> inputShapesMapTo1d = {{1, 1, 56, 5}, {1, 32, 56, 5}, {1, 2, 64, 5}};

const auto conv2DParams_Kernels2D = ::testing::Combine(::testing::ValuesIn(kernels2D),
                                                       ::testing::ValuesIn(strides2D),
                                                       ::testing::ValuesIn(padBegins2D),
                                                       ::testing::ValuesIn(padEnds2D),
                                                       ::testing::ValuesIn(dilations2D),
                                                       ::testing::ValuesIn(numOutCannels2D),
                                                       ::testing::Values(ngraph::op::PadType::EXPLICIT));
const auto conv2DParams_Kernels2D_big = ::testing::Combine(::testing::ValuesIn(kernels2D_big),
                                                           ::testing::ValuesIn(strides2D),
                                                           ::testing::ValuesIn(padBegins2D),
                                                           ::testing::ValuesIn(padEnds2D),
                                                           ::testing::ValuesIn(dilations2D),
                                                           ::testing::ValuesIn(numOutCannels2D),
                                                           ::testing::Values(ngraph::op::PadType::EXPLICIT));

const auto conv2DParams_Kernels2D_3x3 = ::testing::Combine(::testing::ValuesIn(kernels2D_3x3),
                                                           ::testing::ValuesIn(strides2D),
                                                           ::testing::ValuesIn(padBegins2D),
                                                           ::testing::ValuesIn(padEnds2D),
                                                           ::testing::ValuesIn(dilations2D),
                                                           ::testing::ValuesIn(num_out_channels_for_mapped_2d),
                                                           ::testing::Values(ngraph::op::PadType::EXPLICIT));

const auto conv2DParams_Kernels2D_5x6 = ::testing::Combine(::testing::ValuesIn(kernels2D_5x6),
                                                           ::testing::ValuesIn(strides2D),
                                                           ::testing::ValuesIn(padBegins2D),
                                                           ::testing::ValuesIn(padEnds2D),
                                                           ::testing::ValuesIn(dilations2D),
                                                           ::testing::ValuesIn(num_out_channels_for_mapped_2d),
                                                           ::testing::Values(ngraph::op::PadType::EXPLICIT));

const auto conv2DParams_ExplicitPadding_Height1 = ::testing::Combine(::testing::ValuesIn(kernelsH1),
                                                                     ::testing::ValuesIn(stridesH1),
                                                                     ::testing::ValuesIn(padBeginsH1),
                                                                     ::testing::ValuesIn(padEndsH1),
                                                                     ::testing::ValuesIn(dilationsH1),
                                                                     ::testing::ValuesIn(numOutCannels),
                                                                     ::testing::Values(ngraph::op::PadType::EXPLICIT));
const auto conv2DParams_ExplicitPadding_Width1 = ::testing::Combine(::testing::ValuesIn(kernelsW1),
                                                                    ::testing::ValuesIn(stridesW1),
                                                                    ::testing::ValuesIn(padBeginsW1),
                                                                    ::testing::ValuesIn(padEndsW1),
                                                                    ::testing::ValuesIn(dilationsW1),
                                                                    ::testing::ValuesIn(numOutCannels),
                                                                    ::testing::Values(ngraph::op::PadType::EXPLICIT));
const auto conv2DParams_AutoPadValid_Height1 = ::testing::Combine(::testing::ValuesIn(kernelsH1),
                                                                  ::testing::ValuesIn(stridesH1),
                                                                  ::testing::Values(std::vector<ptrdiff_t>({0, 0})),
                                                                  ::testing::Values(std::vector<ptrdiff_t>({0, 0})),
                                                                  ::testing::ValuesIn(dilationsH1),
                                                                  ::testing::ValuesIn(numOutCannels),
                                                                  ::testing::Values(ngraph::op::PadType::VALID));
const auto conv2DParams_AutoPadValid_Height1_BigStride =
    ::testing::Combine(::testing::ValuesIn(kernelsH1),
                       ::testing::ValuesIn(stridesH1Big),
                       ::testing::Values(std::vector<ptrdiff_t>({0, 0})),
                       ::testing::Values(std::vector<ptrdiff_t>({0, 0})),
                       ::testing::ValuesIn(dilationsH1),
                       ::testing::ValuesIn(numOutCannels),
                       ::testing::Values(ngraph::op::PadType::VALID));
const auto conv2DParams_AutoPadValid_Width1 = ::testing::Combine(::testing::ValuesIn(kernelsW1),
                                                                 ::testing::ValuesIn(stridesW1),
                                                                 ::testing::Values(std::vector<ptrdiff_t>({0, 0})),
                                                                 ::testing::Values(std::vector<ptrdiff_t>({0, 0})),
                                                                 ::testing::ValuesIn(dilationsW1),
                                                                 ::testing::ValuesIn(numOutCannels),
                                                                 ::testing::Values(ngraph::op::PadType::VALID));
const auto conv2DParams_AutoPadValid_MapTo1d = ::testing::Combine(::testing::Values(std::vector<size_t>{3, 5}),
                                                                  ::testing::ValuesIn(stridesW1),
                                                                  ::testing::Values(std::vector<ptrdiff_t>({0, 0})),
                                                                  ::testing::Values(std::vector<ptrdiff_t>({0, 0})),
                                                                  ::testing::Values(std::vector<size_t>{1, 1}),
                                                                  ::testing::ValuesIn(numOutCannels),
                                                                  ::testing::Values(ngraph::op::PadType::VALID));

// TODO: padding isn't currently supported in GNA
INSTANTIATE_TEST_SUITE_P(DISABLED_smoke_Convolution2D_ExplicitPadding_Height1,
                         ConvolutionLayerTest,
                         ::testing::Combine(conv2DParams_ExplicitPadding_Height1,
                                            ::testing::ValuesIn(netPrecisions),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::ValuesIn(inputShapesH1),
                                            ::testing::Values(CommonTestUtils::DEVICE_GNA)),
                         ConvolutionLayerTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(DISABLED_smoke_Convolution2D_ExplicitPadding_Width1,
                         ConvolutionLayerTest,
                         ::testing::Combine(conv2DParams_ExplicitPadding_Width1,
                                            ::testing::ValuesIn(netPrecisions),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::ValuesIn(inputShapesW1),
                                            ::testing::Values(CommonTestUtils::DEVICE_GNA)),
                         ConvolutionLayerTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(smoke_Convolution2D_AutoPadValid_Height1,
                         ConvolutionLayerTest,
                         ::testing::Combine(conv2DParams_AutoPadValid_Height1,
                                            ::testing::ValuesIn(netPrecisions),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::ValuesIn(inputShapesH1),
                                            ::testing::Values(CommonTestUtils::DEVICE_GNA)),
                         ConvolutionLayerTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(smoke_Convolution2D_AutoPadValid_Height1_BigStride,
                         ConvolutionLayerTest,
                         ::testing::Combine(conv2DParams_AutoPadValid_Height1_BigStride,
                                            ::testing::ValuesIn(netPrecisions),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::ValuesIn(inputShapesH1),
                                            ::testing::Values(CommonTestUtils::DEVICE_GNA)),
                         ConvolutionLayerTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(smoke_Convolution2D_AutoPadValid_Width1,
                         ConvolutionLayerTest,
                         ::testing::Combine(conv2DParams_AutoPadValid_Width1,
                                            ::testing::ValuesIn(netPrecisions),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::ValuesIn(inputShapesW1),
                                            ::testing::Values(CommonTestUtils::DEVICE_GNA)),
                         ConvolutionLayerTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(smoke_Convolution2D_AutoPadValid_MapTo1d,
                         ConvolutionLayerTest,
                         ::testing::Combine(conv2DParams_AutoPadValid_MapTo1d,
                                            ::testing::ValuesIn(netPrecisions),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::ValuesIn(inputShapesMapTo1d),
                                            ::testing::Values(CommonTestUtils::DEVICE_GNA)),
                         ConvolutionLayerTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(smoke_Convolution2D_Kernels2D,
                         gna_convolution_layerTest,
                         ::testing::Combine(conv2DParams_Kernels2D,
                                            ::testing::ValuesIn(netPrecisions),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::Values(input2DNCHW),
                                            ::testing::Values(CommonTestUtils::DEVICE_GNA)),
                         gna_convolution_layerTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(smoke_Convolution2D_Kernels2D_big,
                         gna_convolution_layerTest,
                         ::testing::Combine(conv2DParams_Kernels2D_big,
                                            ::testing::ValuesIn(netPrecisions),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::Values(input2DNCHW),
                                            ::testing::Values(CommonTestUtils::DEVICE_GNA)),
                         gna_convolution_layerTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(smoke_Convolution2D_Map2D_Not_Transpose_h_w_3_3,
                         ConvolutionLayerTest,
                         ::testing::Combine(conv2DParams_Kernels2D_3x3,
                                            ::testing::ValuesIn(netPrecisions),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::Values(input2DNCHW_3x3),
                                            ::testing::Values(CommonTestUtils::DEVICE_GNA)),
                         ConvolutionLayerTest::getTestCaseName);

INSTANTIATE_TEST_SUITE_P(smoke_Convolution2D_Map2D_Not_Transpose_h_w_5_6,
                         ConvolutionLayerTest,
                         ::testing::Combine(conv2DParams_Kernels2D_5x6,
                                            ::testing::ValuesIn(netPrecisions),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::Values(InferenceEngine::Layout::ANY),
                                            ::testing::Values(input2DNCHW_5x6),
                                            ::testing::Values(CommonTestUtils::DEVICE_GNA)),
                         ConvolutionLayerTest::getTestCaseName);
}  // namespace
